import os
import joblib
from xgboost import XGBClassifier
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from utils.data_load import dataload
from utils.feature_engineering_xgb import feature_engineering

# 1️⃣ 数据加载
X_train, X_test, y_train, y_test = dataload()

# 2️⃣ 特征工程
X_train_scaled, X_test_scaled, y_train_res, encoder, scaler = feature_engineering(X_train, X_test, y_train)

# 3️⃣ 参数搜索网格
param_grid = {
    'n_estimators': [200, 400],
    'max_depth': [3, 5, 7],
    'learning_rate': [0.01, 0.05, 0.1],
    'subsample': [0.7, 0.8, 1.0],
    'colsample_bytree': [0.7, 0.8, 1.0],
}

base_model = XGBClassifier(
    random_state=42,
    eval_metric='logloss',
    scale_pos_weight=(len(y_train_res)-sum(y_train_res))/sum(y_train_res)
)

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

grid_search = GridSearchCV(
    estimator=base_model,
    param_grid=param_grid,
    scoring='roc_auc',
    cv=cv,
    n_jobs=-1,
    verbose=2
)

print("🔍 正在进行超参数搜索...")
grid_search.fit(X_train_scaled, y_train_res)

print("\n✅ 最优参数:", grid_search.best_params_)
print("🌟 最优CV得分:", grid_search.best_score_)

# 4️⃣ 保存模型与编码器
os.makedirs("../model", exist_ok=True)
joblib.dump(grid_search.best_estimator_, "../model/xgboost_model.pkl")

print("✅ 模型、编码器与标准化器已保存！")
